# Helpers for current-flow betweenness and current-flow closness
# Lazy computations for inverse Laplacian and flow-matrix rows.
import networkx as nx

def flow_matrix_row(G, weight='weight', dtype=float, solver='lu'):
    # Generate a row of the current-flow matrix
    import numpy as np
    from scipy import sparse
    from scipy.sparse import linalg
    solvername={"full" :FullInverseLaplacian,
                "lu": SuperLUInverseLaplacian,
                "cg": CGInverseLaplacian}
    n = G.number_of_nodes()
    L = laplacian_sparse_matrix(G, nodelist=range(n), weight=weight, 
                                dtype=dtype, format='csc')
    C = solvername[solver](L, dtype=dtype) # initialize solver
    w = C.w # w is the Laplacian matrix width
    # row-by-row flow matrix
    for u,v,d in G.edges_iter(data=True):
        B = np.zeros(w, dtype=dtype)
        c = d.get(weight,1.0)
        B[u%w] = c
        B[v%w] = -c
        # get only the rows needed in the inverse laplacian 
        # and multiply to get the flow matrix row
        row = np.dot(B, C.get_rows(u,v))  
        yield row,(u,v) 


# Class to compute the inverse laplacian only for specified rows
# Allows computation of the current-flow matrix without storing entire
# inverse laplacian matrix
class InverseLaplacian(object):
    def __init__(self, L, width=None, dtype=None):
        global np
        import numpy as np
        (n,n) = L.shape
        self.dtype = dtype
        self.n = n
        if width is None:
            self.w = self.width(L)
        else:
            self.w = width
        self.C = np.zeros((self.w,n), dtype=dtype)
        self.L1 = L[1:,1:]
        self.init_solver(L)

    def init_solver(self,L):
        pass

    def solve(self,r):
        raise("Implement solver")

    def solve_inverse(self,r):
        raise("Implement solver")


    def get_rows(self, r1, r2):
        for r in range(r1, r2+1):
            self.C[r%self.w, 1:] = self.solve_inverse(r)
        return self.C

    def get_row(self, r):
        self.C[r%self.w, 1:] = self.solve_inverse(r)
        return self.C[r%self.w]


    def width(self,L):
        m=0
        for i,row in enumerate(L):
            w=0
            x,y = np.nonzero(row)
            if len(y) > 0:
                v = y-i
                w=v.max()-v.min()+1
                m = max(w,m)
        return m

class FullInverseLaplacian(InverseLaplacian):
    def init_solver(self,L):
        self.IL = np.zeros(L.shape, dtype=self.dtype)
        self.IL[1:,1:] = np.linalg.inv(self.L1.todense())

    def solve(self,rhs):
        s = np.zeros(rhs.shape, dtype=self.dtype)
        s = np.dot(self.IL,rhs)
        return s

    def solve_inverse(self,r):
        return self.IL[r,1:]


class SuperLUInverseLaplacian(InverseLaplacian):
    def init_solver(self,L):
        from scipy.sparse import linalg
        self.lusolve = linalg.factorized(self.L1.tocsc())

    def solve_inverse(self,r):
        rhs = np.zeros(self.n, dtype=self.dtype)
        rhs[r]=1
        return self.lusolve(rhs[1:])

    def solve(self,rhs):
        s = np.zeros(rhs.shape, dtype=self.dtype)
        s[1:]=self.lusolve(rhs[1:])
        return s



class CGInverseLaplacian(InverseLaplacian):
    def init_solver(self,L):
        global linalg
        from scipy.sparse import linalg
        ilu= linalg.spilu(self.L1.tocsc())
        n=self.n-1
        self.M = linalg.LinearOperator(shape=(n,n), matvec=ilu.solve)

    def solve(self,rhs):
        s = np.zeros(rhs.shape, dtype=self.dtype)
        s[1:]=linalg.cg(self.L1, rhs[1:], M=self.M)[0]
        return s

    def solve_inverse(self,r):
        rhs = np.zeros(self.n, self.dtype)
        rhs[r] = 1
        return linalg.cg(self.L1, rhs[1:], M=self.M)[0]


# graph laplacian, sparse version, will move to linalg/laplacianmatrix.py
def laplacian_sparse_matrix(G, nodelist=None, weight='weight', dtype=None,
                            format='csr'):
    import numpy as np
    import scipy.sparse
    A = nx.to_scipy_sparse_matrix(G, nodelist=nodelist, weight=weight, 
                                  dtype=dtype, format=format)
    (n,n) = A.shape
    data = np.asarray(A.sum(axis=1).T)
    D = scipy.sparse.spdiags(data,0,n,n, format=format)
    return  D - A
